import torch
import torch.nn as nn
import pytorch_lightning as pl
from torch.utils.data import Dataset, DataLoader


class PatchEmbedding(nn.Module):
    """增强稳定性的片段嵌入模块，含动态长度校验"""

    def __init__(self, input_dim, patch_size, embed_dim):
        super().__init__()
        self.patch_size = patch_size
        self.projection = nn.Linear(input_dim * patch_size, embed_dim)
        self.norm = nn.LayerNorm(embed_dim)  # 稳定训练过程

    def forward(self, x):
        batch_size, seq_len, feat_dim = x.size()
        # 动态调整序列长度确保可被patch_size整除
        num_patches = seq_len // self.patch_size
        valid_len = num_patches * self.patch_size
        x = x[:, :valid_len, :]  # 自动截断

        x = x.view(batch_size, num_patches, feat_dim * self.patch_size)
        x = self.projection(x)
        return self.norm(x)


class PositionalEncoding(nn.Module):
    """增强型可扩展位置编码"""

    def __init__(self, embed_dim, max_len=5000):
        super().__init__()
        self.embed_dim = embed_dim
        self.register_buffer('pe', self._create_pe(max_len, embed_dim))

    def _create_pe(self, max_len, embed_dim):
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, embed_dim, 2) * (-torch.log(torch.tensor(10000.0)) / embed_dim))
        pe = torch.zeros(1, max_len, embed_dim)
        pe[0, :, 0::2] = torch.sin(position * div_term)
        pe[0, :, 1::2] = torch.cos(position * div_term)
        return pe

    def forward(self, x):
        return x + self.pe[:, :x.size(1), :]


class PathFormer(pl.LightningModule):
    """与项目架构深度集成的时空预测模型"""

    def __init__(self, input_dim, patch_size=30, embed_dim=512,
                 num_heads=16, num_layers=8, output_dim=16,
                 dropout=0.2, ff_dim=2048):
        super().__init__()
        self.save_hyperparameters()  # 保存超参数便于加载

        # 输入处理模块
        self.patch_embed = PatchEmbedding(input_dim, patch_size, embed_dim)
        self.pos_encoder = PositionalEncoding(embed_dim)

        # Transformer核心模块
        self.transformer = nn.TransformerEncoder(
            encoder_layer=nn.TransformerEncoderLayer(
                d_model=embed_dim,
                nhead=num_heads,
                dim_feedforward=ff_dim,  # 可配置前馈维度
                dropout=dropout,
                batch_first=True,
                activation='gelu'  # 更平滑的激活函数
            ),
            num_layers=num_layers
        )

        # 输出模块
        self.dropout = nn.Dropout(dropout)
        self.fc = nn.Linear(embed_dim, output_dim)

        # 初始化参数
        self._init_weights()

    def _init_weights(self):
        """Xavier初始化提升收敛稳定性"""
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.xavier_normal_(module.weight)
                if module.bias is not None:
                    nn.init.constant_(module.bias, 0)

    def forward(self, x):
        x = self.patch_embed(x)  # (B, L, D)
        x = self.pos_encoder(x)  # 注入位置信息
        x = self.transformer(x)  # (B, L, D)
        x = x.mean(dim=1)  # 时序聚合
        return self.fc(self.dropout(x))

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = nn.HuberLoss(delta=1.0)(y_hat, y)  # 鲁棒性损失
        self.log('train_loss', loss, prog_bar=True)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        val_loss = nn.HuberLoss(delta=1.0)(y_hat, y)
        self.log('val_loss', val_loss, prog_bar=True)

    def configure_optimizers(self):
        """优化器配置与run.py中的训练设置完全兼容"""
        optimizer = torch.optim.AdamW(
            self.parameters(),
            lr=1e-4,
            weight_decay=0.01
        )
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
            optimizer,
            T_max=300,
            eta_min=1e-6
        )
        return [optimizer], [scheduler]


class TimeSeriesDataset(Dataset):
    """与time_sequence.py完全兼容的数据集类"""

    def __init__(self, X, y):
        self.X = torch.as_tensor(X, dtype=torch.float32)
        self.y = torch.as_tensor(y, dtype=torch.float32)

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]


def create_dataloaders(X, y, batch_size=128, workers=4):
    """与run.py训练流程兼容的数据加载器"""
    dataset = TimeSeriesDataset(X, y)
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=workers,
        persistent_workers=workers > 0,
        pin_memory=True  # 加速GPU传输
    )
